import gradio as gr
from pinecone import Pinecone
import numpy as np
from skimage.transform import resize
import matplotlib.pyplot as plt
from collections import Counter

# Pinecone 配置
PINECONE_API_KEY = "pcsk_3Du3r4_PPKjDJ3RXs6svx1MWTwTgVEK3VSpdFekSf1aH5TqmFqzzYoSUUCxhRU2ZRXQ1jQ"
INDEX_NAME = "mnist-index"

# 8x8 MNIST
_side = 8
_scale_max = 16

# 初始化 Pinecone 客户端
pinecone = Pinecone(api_key=PINECONE_API_KEY)
index = pinecone.Index(INDEX_NAME)
print(f"[INFO] Pinecone index '{INDEX_NAME}' 已连接")

def predict_digit(image, invert=False):
    if image is None:
        return "请在画布上绘制一个数字。"
    try:
        if isinstance(image, dict):
            if 'image' in image and image['image'] is not None:
                image = image['image']
            elif 'composite' in image and image['composite'] is not None:
                image = image['composite']
            else:
                return "请在画布上绘制一个数字。"
        if not isinstance(image, np.ndarray):
            image = np.array(image)
    except Exception as e:
        return f"无法解析输入图像: {e}"
    if not isinstance(image, np.ndarray) or image.size == 0 or getattr(image, 'ndim', 0) < 2:
        return "请在画布上绘制一个更清晰的数字。"
    # 灰度化
    if image.ndim == 3:
        if image.shape[2] == 4:
            image_gray = image[:, :, :3].mean(axis=2)
        else:
            image_gray = image.mean(axis=2)
    else:
        image_gray = image
    # 反色（可选）
    if invert:
        if image_gray.max() > 1.0:
            image_gray = image_gray / 255.0
        image_gray = 1.0 - image_gray
    # 归一化到 0..1
    try:
        if image_gray.max() > 1.0:
            image_gray = image_gray / 255.0
    except Exception:
        return "请在画布上绘制一个数字（当前图像无效）。"
    # 自动反色
    try:
        if not invert and float(image_gray.mean()) > 0.5:
            image_gray = 1.0 - image_gray
    except Exception:
        pass
    # 裁剪前景、居中、缩放
    try:
        thr = max(0.2, float(image_gray.mean()) * 0.7)
        mask = image_gray > thr
        if mask.sum() < 10:
            return "笔迹太淡或太小，请加粗并画大一些。"
        rows = np.any(mask, axis=1)
        cols = np.any(mask, axis=0)
        rmin, rmax = np.where(rows)[0][[0, -1]]
        cmin, cmax = np.where(cols)[0][[0, -1]]
        cropped = image_gray[rmin:rmax+1, cmin:cmax+1]
        h, w = cropped.shape
        L = max(h, w)
        pad_vert = L - h
        pad_horz = L - w
        top = pad_vert // 2
        bottom = pad_vert - top
        left = pad_horz // 2
        right = pad_horz - left
        squared = np.pad(cropped, ((top, bottom), (left, right)), mode='constant', constant_values=0.0)
        image_resized = resize(squared, (_side, _side), anti_aliasing=True)
    except Exception:
        image_resized = resize(image_gray, (_side, _side), anti_aliasing=True)
    image_resized = np.clip(image_resized * _scale_max, 0, _scale_max)
    image_resized = image_resized.astype(int)
    image_flattened = image_resized.flatten().tolist()
    # Pinecone 查询
    try:
        results = index.query(
            vector=image_flattened,
            top_k=11,
            include_metadata=True
        )
        labels = [match['metadata']['label'] for match in results['matches']]
        if not labels:
            return "未找到匹配结果，请检查索引数据。"
        pred = Counter(labels).most_common(1)[0][0]
    except Exception as e:
        return f"Pinecone 查询失败: {e}"
    try:
        plt.imsave('processed_image.png', image_resized, cmap='gray')
    except Exception:
        pass
    return f"{pred}"


with gr.Blocks(theme="dark") as demo:
    gr.Markdown("# 手写数字识别（手绘版）")
    with gr.Row():
        try:
            # 不传入额外参数以保证最大兼容性
            image_input = gr.Sketchpad(label="image")
            hint_txt = gr.Markdown("*提示：如果未出现手写画布，请刷新页面（Ctrl+F5）或升级 gradio。*")
        except Exception:
            image_input = gr.Image(image_mode="L", label="image")
            hint_txt = gr.Markdown("*当前环境不支持手写画布，已回退为图片上传。可运行 `pip install --upgrade gradio` 升级。*")
        output_box = gr.Textbox(label="output", max_lines=1)
    invert_checkbox = gr.Checkbox(label="反色（将背景变黑，笔迹变白）", value=False)
    with gr.Row():
        submit_btn = gr.Button("识别")
        clear_btn = gr.Button("清除")

    def clear_canvas():
        return None, ""

    submit_btn.click(predict_digit, inputs=[image_input, invert_checkbox], outputs=output_box)
    clear_btn.click(clear_canvas, inputs=None, outputs=[image_input, output_box])

demo.launch()
